import { h, Service, SessionError } from "koishi"
import { readFile } from "fs/promises"
import { fileURLToPath } from "url"
import { Worker } from "worker_threads"

const workerFunc = () => {
  const { parentPort } = require("worker_threads")
  const uiua = require("@dgck81lnn/uiua")
  parentPort.on("message", ({ action, args, nonce }) => {
    try {
      const result = uiua[action](...args)
      parentPort.postMessage({ nonce, success: true, result })
    } catch (error) {
      parentPort.postMessage({ nonce, success: false, error })
    }
  })
}

export class UiuaError extends Error {
  name = "UiuaError"
}

export class UiuaService extends Service {
  _workerReady
  _worker
  _counter = 1
  _callbacks = new Map()

  static name = "uiua"
  static inject = { optional: ["silk", "ffmpeg"] }

  _setupWorker() {
    this.ctx.logger.info("starting worker")
    this._worker = new Worker(`(${workerFunc})()`, {
      name: "uiua-worker",
      eval: true,
    })
    this._worker.on("message", ({ nonce, success, result, error }) => {
      if (!this._callbacks.has(nonce)) return
      this._callbacks.get(nonce)(success ? result : Promise.reject(error))
      this._callbacks.delete(nonce)
    })
    return new Promise(res => this._worker.on("online", res))
  }

  async _run(action, ...args) {
    const nonce = this._counter++
    await (this._workerReady ??= this._setupWorker())
    this._worker.postMessage({ action, args, nonce })
    return new Promise(res => {
      const clearTO = this.ctx.setTimeout(() => {
        this._callbacks.get(nonce)(Promise.reject(new UiuaError("worker-timeout")))
        this._callbacks.delete(nonce)
      }, 10000)
      this._callbacks.set(nonce, result => {
        clearTO()
        res(result)
      })
    })
  }

  eval(...args) {
    return this._run("eval", ...args)
  }
  eval_mm(...args) {
    return this._run("eval_mm", ...args)
  }
  format(...args) {
    return this._run("format", ...args)
  }

  async readInput(type, attrs) {
    if (!attrs?.src) return
    try {
      let data = /[/\\]/.test(attrs.path)
        ? await readFile(attrs.path)
        : attrs.src.startsWith("file:")
        ? await readFile(fileURLToPath(attrs.src))
        : attrs.src.startsWith("base64://")
        ? Buffer.from(attrs.src.slice(9), "base64")
        : await this.ctx.http.get(attrs.src, { responseType: "arraybuffer" })

      this.ctx.logger.debug("data", data)

      if (type === "audio") {
        let pcmBytes
        if (this.ctx.silk && this.ctx.silk.isSilk(data)) {
          pcmBytes = (await this.ctx.silk.decode(data, 44100)).data
        } else if (this.ctx.ffmpeg) {
          const ff = this.ctx.ffmpeg.builder()
          ff.input(data instanceof Buffer ? data : Buffer.from(data))
          ff.outputOption("-f", "s16le")
          pcmBytes = await ff.run("buffer")
        }
        if (pcmBytes) {
          // Prepend WAV header
          const buffer = Buffer.alloc(pcmBytes.length + 44)
          ;[
            0x46464952, // RIFF
            pcmBytes.length + 36,
            0x45564157, // WAVE
            0x20746d66, // fmt
            16,
            0x00010001,
            44100,
            88200,
            0x00100002,
            0x61746164, // data
            pcmBytes.length,
          ].forEach((v, i) => {
            buffer.writeUInt32LE(v, i << 2)
          })
          buffer.set(pcmBytes, 44)
          this.ctx.logger.debug("audio transformed", buffer)
          return buffer
        }
      }

      return data instanceof Buffer ? data : Buffer.from(data)
    } catch (err) {
      this.ctx.logger.error(err)
      throw new UiuaError("failed-to-read-input")
    }
  }

  stop() {
    this.ctx.logger.info("terminating worker")
    this._worker?.terminate()
  }

  constructor(ctx) {
    super(ctx, "uiua")

    ctx.command("uiua <code:rawtext>", { strictOptions: true, captureQuote: false })
      .option("format", "-f")
      .option("format", "-F, --format-only", { value: "only" })
      .option("input", "-I <fragment:fragment>", { hidden: true })
      .option("image", "-i <image:image>", { hidden: true })
      .option("audio", "-a <audio:audio>", { hidden: true })
      .option("rest", "-- <code:rawtext>", { descPath: "lnnbot.general.rest-option" })
      .action(async ({ root, session, options }, code) => {
        const { format } = options
        const formatOnly = format === "only"
        code ||= options.rest || ""
        let input =
          formatOnly ||
          (await this.readInput("img", options.image)) ||
          (await this.readInput("audio", options.audio))
        const inputEls = options.input || (root === true && session.quote?.elements)
        if (inputEls) {
          const el = inputEls[0]
          if (["img", "image", "audio", "file"].includes(el?.type))
            input ||= await this.readInput(el.type, el.attrs)
          else if (!code) code = h("", inputEls).toString(true)
          else input ||= h("", inputEls).toString(true)
        }
        if (!code && (formatOnly || !input)) return h.i18n("internal.insufficient-arguments")

        if (options.format)
          try {
            code = await this.format(code)
          } catch (err) {
            ctx.logger.debug(err)
            return h.text(err.toString())
          }
        if (formatOnly) return h.text(code)

        try {
          let { outputs, error } = await this.eval_mm(code, {
            inputs: input ? [input] : [],
            experimental: true,
            timeout: 8000,
          })
          let prevType
          const result = []
          if (format) result.push(h.text(code), h("message"))
          if (error) outputs = outputs.concat([{ type: "error", content: error }])
          if (outputs.length) {
            for (const item of outputs) {
              if (prevType && prevType !== item.type) {
                const prev = result.at(-1)
                if (typeof prev === "string")
                  prev.endsWith("\n") || (result[result.length - 1] += "\n")
                else result.push("\n")
              }
              if (["image", "audio"].includes(item.type)) {
                if (item.label) result.push(item.label + ": ")
                result.push(h[item.type](item.data, item.mime))
              } else {
                result.push(item.content)
              }
              prevType = item.type
            }
          } else {
            result.push(...session.i18n(".no-output"))
          }
          return result
        } catch (err) {
          if (err instanceof UiuaError) throw new SessionError("." + err.message)
          ctx.logger.debug(err)
          return h.text(err.toString())
        }
      })

    ctx.middleware((session, next) => {
      if (session.stripped.hasAt && !session.stripped.atSelf) return next()
      const prefix = session.stripped.content.match(/^∪+/)?.[0]
      if (prefix) {
        if (prefix.length > 3) return next()//prefix + "<censor>" + session.stripped.content.slice(prefix.length) + "</censor>"
        const code = h("", h.parse(session.stripped.content.slice(prefix.length).trimStart())).toString(true)
        if (!code) return next()
        return session.execute({
          name: "uiua",
          args: [code],
          options: prefix === "∪∪∪" ? { format: "only" } : prefix === "∪∪" ? { format: true } : {},
          root: true,
        })
      }
      return next()
    })

    ctx.i18n.define("zh-CN", "commands.uiua", {
      description: "运行 Uiua 代码",
      usage: "快捷方式：\n∪(code...) = uiua (code...)\n∪∪(code...) = uiua -f (code...)\n∪∪∪(code...) = uiua -F (code...)",
      options: {
        format: "格式化代码",
        "format.only": "仅格式化代码，不运行",
        input: "输入",
        image: "输入图片",
        audio: "输入音频",
      },
      messages: {
        "worker-timeout": "子线程回应超时…",
        "failed-to-read-input": "读取输入失败…",
        "no-output": "（无输出…）",
      },
    })
  }
}

export default UiuaService
